import sqlite3
import hashlib
import os
from datetime import datetime
from typing import Optional

from src.utils.logger import get_logger

logger = get_logger(__name__)

class TranslationCache:
    """
    A class to cache translation results in a local SQLite database
    to avoid redundant API calls.
    """

    def __init__(self, db_path: Optional[str] = None):
        """
        Initializes the cache and connects to the database.
        Creates the necessary table if it doesn't exist.

        :param db_path: Optional path to the database file. If None, a default
                        path in the user's home directory is used.
        """
        if db_path is None:
            # Store the database in a dedicated folder in the user's home directory
            app_dir = os.path.expanduser("~/.trans-video-app")
            os.makedirs(app_dir, exist_ok=True)
            self.db_path = os.path.join(app_dir, "translation_cache.db")
        else:
            self.db_path = db_path

        self.conn = None
        try:
            self.conn = sqlite3.connect(self.db_path)
            self._create_table()
            self._migrate_table()
        except sqlite3.Error as e:
            logger.error(f"Database error: {e}")
            self.conn = None

    def _create_table(self):
        """Creates the 'translation_cache' table if it doesn't exist."""
        if not self.conn:
            return
        try:
            cursor = self.conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS translation_cache (
                    source_text_hash TEXT NOT NULL,
                    translator_name TEXT NOT NULL,
                    model_name TEXT NOT NULL,
                    source_language TEXT NOT NULL,
                    target_language TEXT NOT NULL,
                    translated_text TEXT NOT NULL,
                    source_text TEXT NOT NULL,
                    created_at TIMESTAMP NOT NULL,
                    PRIMARY KEY (source_text_hash, translator_name, model_name, source_language, target_language)
                )
            """)
            self.conn.commit()
        except sqlite3.Error as e:
            logger.error(f"Failed to create table: {e}")

    def _migrate_table(self):
        """Adds the model_name column if it doesn't exist for backward compatibility."""
        if not self.conn:
            return
        try:
            cursor = self.conn.cursor()
            cursor.execute("PRAGMA table_info(translation_cache)")
            columns = [info[1] for info in cursor.fetchall()]
            if 'model_name' not in columns:
                logger.info("Updating database schema: Adding 'model_name' column to cache.")
                # We have to recreate the table to add a new primary key.
                # This is a simple migration that might lose old data if not handled carefully.
                # For this application, we'll assume losing cache on schema change is acceptable.
                cursor.execute("DROP TABLE IF EXISTS translation_cache")
                self._create_table()
                logger.info("Cache table has been recreated with the new schema.")
            self.conn.commit()
        except sqlite3.Error as e:
            logger.error(f"Failed to migrate table: {e}")

    @staticmethod
    def _hash_text(text: str) -> str:
        """Generates a SHA-256 hash for the given text."""
        return hashlib.sha256(text.encode('utf-8')).hexdigest()

    def get(self, source_text: str, translator_name: str, model_name: str, source_language: str, target_language: str) -> Optional[str]:
        """
        Retrieves a translated text from the cache.

        :param source_text: The original text.
        :param translator_name: The name of the translation service (e.g., 'openai').
        :param model_name: The specific model used for translation (e.g., 'gpt-4').
        :param source_language: The source language code (e.g., 'en').
        :param target_language: The target language code (e.g., 'zh-CN').
        :return: The translated text if found in cache, otherwise None.
        """
        if not self.conn:
            return None

        source_text_hash = self._hash_text(source_text)
        try:
            cursor = self.conn.cursor()
            cursor.execute("""
                SELECT translated_text FROM translation_cache
                WHERE source_text_hash = ? AND translator_name = ? AND model_name = ? AND source_language = ? AND target_language = ?
            """, (source_text_hash, translator_name, model_name, source_language, target_language))
            
            result = cursor.fetchone()
            
            if result:
                logger.info(f"Cache hit for '{source_text[:30]}...'")
                return result[0]
            else:
                logger.info(f"Cache miss for '{source_text[:30]}...'")
                return None
        except sqlite3.Error as e:
            logger.error(f"Failed to get cache: {e}")
            return None

    def set(self, source_text: str, translated_text: str, translator_name: str, model_name: str, source_language: str, target_language: str):
        """
        Stores a new translation in the cache.

        :param source_text: The original text.
        :param translated_text: The translated text.
        :param translator_name: The name of the translation service.
        :param model_name: The specific model used for translation.
        :param source_language: The source language code.
        :param target_language: The target language code.
        """
        if not self.conn:
            return

        source_text_hash = self._hash_text(source_text)
        created_at = datetime.utcnow()
        
        try:
            cursor = self.conn.cursor()
            cursor.execute("""
                INSERT OR REPLACE INTO translation_cache
                (source_text_hash, translator_name, model_name, source_language, target_language, translated_text, source_text, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (source_text_hash, translator_name, model_name, source_language, target_language, translated_text, source_text, created_at))
            self.conn.commit()
            logger.info(f"Cached translation for '{source_text[:30]}...'")
        except sqlite3.Error as e:
            logger.error(f"Failed to set cache: {e}")

    def close(self):
        """Closes the database connection."""
        if self.conn:
            self.conn.close()
            self.conn = None

# Example usage:
if __name__ == '__main__':
    cache = TranslationCache(db_path=':memory:') # Use in-memory DB for example
    
    # --- Test Case 1: Cache miss ---
    print("--- Test Case 1 ---")
    retrieved = cache.get("Hello, world", "openai", "gpt-4", "en", "zh-CN")
    print(f"Retrieved: {retrieved}") # Expected: None
    
    # --- Test Case 2: Cache set ---
    print("\n--- Test Case 2 ---")
    cache.set("Hello, world", "你好，世界", "openai", "gpt-4", "en", "zh-CN")
    
    # --- Test Case 3: Cache hit ---
    print("\n--- Test Case 3 ---")
    retrieved = cache.get("Hello, world", "openai", "gpt-4", "en", "zh-CN")
    print(f"Retrieved: {retrieved}") # Expected: 你好，世界

    # --- Test Case 4: Different model (cache miss) ---
    print("\n--- Test Case 4 ---")
    retrieved = cache.get("Hello, world", "openai", "gpt-3.5-turbo", "en", "zh-CN")
    print(f"Retrieved: {retrieved}") # Expected: None

    cache.close()